import asyncio
import websockets as ws
from websockets import ConnectionClosed


async def send_msg(websocket):
    while True:
        try:
            await input_cmd(websocket)
            recv = await websocket.recv()
            print(f"recv: len{len(recv)}")
        except ConnectionClosed as e:
            print(e.code)
            if e.code == 1006:
                print("restart")
                await asyncio.sleep(2)
                break


async def webscoket_client():
    uri = "ws://localhost:8899"
    while True:
        try:
            async with ws.connect(uri) as websocket:
                await send_msg(websocket)

        except ConnectionRefusedError as e:
            print(e)
            await asyncio.sleep(2)


async def input_cmd(websocket):
    help_str = """
输入命令, 指定发送数据大小
1 512B, 2 1024B, 3 2048B, 4 4096B
q 退出

"""
    data = b'12345678'
    cmd = input(help_str)
    if cmd == "1":
        multi = 64
    elif cmd == "2":
        multi = 128
    elif cmd == "3":
        multi = 256
    elif cmd == "4":
        multi = 512
    else:
        print("error 未知指令")
        print(help_str)

    await websocket.send(data * multi)


if __name__ == "__main__":
    # asyncio.run(webscoket_client) --> NG
    asyncio.get_event_loop().run_until_complete(webscoket_client())
